# MIT License
# 
# Copyright (c) 2016 David Sandberg
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

#
# Borrowed from davidsandberg's facenet project: https://github.com/davidsandberg/facenet
# From this directory:
#   facenet/src/align
#
# Just keep the MTCNN related stuff and removed other codes
# python package required:
#     tensorflow, opencv,numpy\

# -*- coding: utf-8 -*-


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import sys
import os
import os.path
import argparse
import tensorflow as tf
import numpy as np
from mtcnn.mtcnn import mtcnn
import cv2
import time

def zxzhao_mkdirs(my_path):
	if not os.path.exists(my_path):
		os.makedirs(my_path)
def save_img(img, img_name,save_path):
	zxzhao_mkdirs(save_path)
	save_path2 = os.path.join(save_path,img_name)
	cv2.imwrite(save_path2, img)

def file_extension(path,index): 
    return os.path.splitext(path)[index] 

def main(args):
    inputDir=args.input
    print ("dir:"+inputDir)

    m_mtcnn = mtcnn()

    controled_size = 256

    img_num = 0
    time_start=time.time()
    ######################################
    for rootpath,dirnames,filenames in os.walk(inputDir):
        for cur_filename in filenames:
            filename=os.path.join(rootpath, cur_filename)
            if file_extension(cur_filename,1) != ".jpg":
                continue
    ######################################
            print ("file:"+filename)
            
            try:
                draw = cv2.imread(filename)

                # remove error image
                img_size = draw.shape
                height = img_size[0]#height(rows) of image
                width = img_size[1]#width(colums) of image
                if height< 20 or width< 20:
                    os.remove(filename)
                    continue

                bounding_boxes, points = m_mtcnn.detect_face_v3(filename)
                # print (bounding_boxes, points)
                # print (type(bounding_boxes))
                # print (type(points))

                # continue
                if type(bounding_boxes) is np.ndarray:
                    nrof_faces = bounding_boxes.shape[0]
                    print('Total %d face(s) detected' % (nrof_faces))
                    if nrof_faces==1: 
                        # continue                 
                        for rect in bounding_boxes:
                            # margin_w = (int(rect[3]) - int(rect[1]))/5
                            # margin_h = (int(rect[2]) - int(rect[0]))/5
                            margin_w = (int(rect[3]) - int(rect[1]))/7
                            margin_h = (int(rect[2]) - int(rect[0]))/7
                            bb = np.zeros(4, dtype=np.int32)
                            bb[0] = np.maximum(rect[0] - margin_w, 0)
                            bb[1] = np.maximum(rect[1] - margin_h, 0)
                            bb[2] = np.minimum(rect[2] + margin_w, img_size[1])
                            bb[3] = np.minimum(rect[3] + margin_h, img_size[0])
                            print (bb)

                            if (int(bb[2])-int(bb[0]))<50 or (int(bb[3])-int(bb[1]))<50:
                                continue
                            else:
                                img_num +=1
                                face_crop_img = draw[bb[1]:bb[3],bb[0]:bb[2]]
                                img_height = bb[3]-bb[1]
                                img_width  = bb[2]-bb[0]

                                min_size = min(img_height, img_width)
                                resize_ratio = float(controled_size)/float(min_size)
                                if resize_ratio > 1.0 :
                                	resize_ratio = 1.0
                                else:
                                	face_crop_img = cv2.resize(face_crop_img, (int(img_width*resize_ratio), int(img_height*resize_ratio)))

                            cv2.imwrite(filename, face_crop_img)
                    else:
                        os.remove(filename)
                else:
                    os.remove(filename)
            except:
                continue
    print ("face_num = " + str(img_num))
            

def parse_arguments(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str, help='image path to be detected for faces.',default='./img/')
    return parser.parse_args(argv)

if __name__ == '__main__':
    main(parse_arguments(sys.argv[1:]))
